#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Thu Feb  9 18:58:28 2017

@author: kevin
"""
import os
import glob
import h5py
import numpy as np
import cv2
import tensorflow as tf
from vgg16 import vgg16

def read_data(path):
     images = np.array(h5py.File(path, 'r').get('images'))
     labels = np.array(h5py.File(path,'r').get('labels'))
     return images, labels
     
def compare(x,y):
     stat_x = os.path.basename(x)
     stat_y = os.path.basename(y)
     if stat_x < stat_y:
          return -1
     elif stat_x > stat_y:
          return 1
     else:
          return 0

def prepare_data(sess, dataset):
     images_dir = os.path.join(os.getcwd(), dataset, 'images')
     labels_dir = os.path.join(os.getcwd(), dataset, 'labels')
     images = glob.glob(os.path.join(images_dir, '*.png'))
     labels = glob.glob(os.path.join(labels_dir, '*.png'))
     images.sort(compare)
     labels.sort(compare)
     return images, labels
     
def imread(path, is_singleChannel, img_sz):
     if is_singleChannel:
          im = cv2.imread(path)
          im = cv2.resize(im, (img_sz, img_sz))
          im_single = im[:,:,1]
          return im_single
     else:
          im = cv2.imread(path)
          im = cv2.resize(im, (img_sz, img_sz))
          return im

def make_data(sess, images, labels):
     savepath = os.path.join(os.getcwd(), 'train.h5')
     with h5py.File(savepath, 'w') as hf:
          hf.create_dataset('images', data=images)
          hf.create_dataset('labels', data=labels)
     
def input_setup(sess, img_sz):
     images, labels = prepare_data(sess, dataset='Train')
     input_seq = []
     label_seq = []

     for i in xrange(len(images)):
          image_ = imread(images[i], False, img_sz=img_sz)
          label_ = imread(labels[i], True, img_sz=img_sz)
          input_seq.append(image_)
          label_seq.append(label_)

     arrimages = np.asarray(input_seq)
     arrlabels = np.asarray(label_seq)
     make_data(sess, arrimages, arrlabels)
        


if __name__=='__main__':
     with tf.Session() as sess:
          input_setup(sess,400)
          data_dir = os.path.join(os.getcwd(), 'train.h5')
          train_images, train_labels = read_data(data_dir)
          print train_images.shape
          print train_labels.shape
          imgs = tf.placeholder(tf.float32, [None,224,224,3])
          vgg = vgg16(imgs, 'vgg16_weights.npz',sess)
          train_probs = []
          for i in range(train_images.shape[0]):
               print i
               img = train_images[i,:,:,:]
               resized_img = cv2.resize(img,(224,224))
               prob = sess.run(vgg.probs, feed_dict={vgg.imgs:[resized_img]})
               train_probs.append(prob)
          savepath = os.path.join(os.getcwd(), 'probs.h5')
          with h5py.File(savepath, 'w') as hf:
                hf.create_dataset('probs', data=train_probs)
     
     
     
     
